"""
dijkstra 单源最短路径算法  python实现
使用优先队列加速

实现
1. 将源点加入堆，并调整堆。 
2. 选出堆顶元素u（即代价最小的元素），从堆中删除，并对堆进行调整。 
3. 处理与u相邻的，未被访问过的，满足三角不等式的顶点 
    1):若该点在堆里，更新距离，并调整该元素在堆中的位置。 
    2):若该点不在堆里，加入堆，更新堆。 
4. 若取到的u为终点，结束算法；否则重复步骤2、3。 

3.1)中的“调整该元素在堆中的位置”并不是修改堆中元素的数值，而是把更新后的数值重新加入堆中，这样这个节点更新后的数值（小），自然在原来数值之前
而堆中同时可以存在一个节点原数值以及多次更新的数值，但是这些冗余都不会产生什么影响，因为如果节点已经确定了（visited），那么这些旧数值只是最后简单弹出而已

堆的主要操作：
1) 构建堆 （heapify 将列表变为堆）
2) 弹出堆顶 （最大值 / 最小值 ， 并且自动调整堆）
3) 将一个元素加入堆 （自动调整堆）
并没有直接修改堆内某个元素的方法，堆heap在抽象设计上也不是一个可以随机访问的数据结构，能接触的只有堆顶
"""
import heapq
import collections

class Solution:

    def dijkstra_all(self, edges, src):
        """
        计算以src为源点到其他所有点的最小距离
        """
        graph = collections.defaultdict(lambda: dict())
        nodes = set()
        for edge in edges:
            begin, end, cost = edge
            graph[begin][end] = cost
            nodes.add(begin)
            nodes.add(end)

        visited = [False] * len(nodes) # 已经确定为最小值 才设置 visited == True
        distance = {n:-1 for n in nodes}

        minHeap = []
        heapq.heappush(minHeap, (0, src)) # 注意 heapq默认使用tuple第一个数值(tuple[0])进行比较建堆

        while minHeap:

            cur_cost, cur_node = heapq.heappop(minHeap)
            
            if visited[cur_node] == False: # 这个节点没有被确定最小距离
                visited[cur_node] = True
                distance[cur_node] = cur_cost

                for dest, cost in graph[cur_node].items(): # 通过这个节点，更新与之相邻的，但是仍未被确定最小距离的节点
                    if visited[dest] == False:
                        heapq.heappush(minHeap, (cur_cost+cost, dest))
        
        return distance

    def dijkstra_src2dst(self, edges, src, dst):
        """
        返回 源点src到目标点dst的距离 如果到达不了 则返回-1
        """

        graph = collections.defaultdict(lambda: dict())
        nodes = set()
        for edge in edges:
            begin, end, cost = edge
            graph[begin][end] = cost
            nodes.add(begin)
            nodes.add(end)

        visited = [False] * len(nodes) # 已经确定为最小值 才设置 visited == True
        
        minHeap = []
        heapq.heappush(minHeap, (0, src)) # 注意 heapq默认使用tuple第一个数值(tuple[0])进行比较建堆

        while minHeap:

            cur_cost, cur_node = heapq.heappop(minHeap)

            if cur_node == dst:
                return cur_cost
            
            if visited[cur_node] == False: # 这个节点没有被确定最小距离
                visited[cur_node] = True

                for dest, cost in graph[cur_node].items(): # 通过这个节点，更新与之相邻的，但是仍未被确定最小距离的节点
                    if visited[dest] == False:
                        heapq.heappush(minHeap, (cur_cost+cost, dest))
        
        return -1




if __name__ == '__main__':
    S = Solution()
    edges, src, dst = [[0,1,1],[0,2,5],[1,2,1],[2,3,1]], 0, 3
    ans = S.dijkstra_src2dst(edges, src, dst)
    print(ans)
    distance = S.dijkstra_all(edges, src)
    print(distance)